Repository navigation
[Feature] Add LeRobot replay dataset adapter - #4611
Conversation
🔗 Helpful Links🧪 See artifacts and rendered test results at hud.pytorch.org/pr/pytorch/rl/4611
Note: Links to docs will display an error until the docs builds have been completed. ✅ No FailuresAs of commit d257b67 with merge base c0ca870 ( This comment was automatically generated by Dr. CI and updates every 15 minutes. |
meta-cla
Bot
added
the
CLA Signed
github-actions
Bot
added
Documentation
theap06
left a comment
theap06
left a comment
There was a problem hiding this comment.
Thanks a bunch! I agree with the overall direction and tested out the logic and wrote to hf. left a couple small comments.
| assert decoded.dtype == torch.uint8 | ||
|
|
||
|
|
||
| @pytest.mark.skipif(not _has_lerobot_dataset, reason="requires torchrl[lerobot]") |
There was a problem hiding this comment.
Can we add a py3.12 stable-torch job to have the ci job run for this?
| fps: int | None = None, | ||
| features: Mapping[str, dict[str, Any]] | None = None, | ||
| key_map: Mapping[str, NestedKey] | None = None, | ||
| episode_key: NestedKey = "episode", |
There was a problem hiding this comment.
right now the episode_key and the step_key duplicate what key_map has. so if we use a keymap to rename we would get a KeyError. can we derive the keys from key_map and drop these args?
| ) | ||
| if isinstance(value, MetaData): | ||
| columns[name] = [value.data] * len(batch) | ||
| elif isinstance(value, torch.Tensor): |
There was a problem hiding this comment.
can we permute chw frames to the feature layout on export? my concern here is that if we import a hub dataset, it gives us chw float images and the the hub schema is hwc.
There was a problem hiding this comment.
Done. I used a simple shape-based heuristic which should work in almost all cases, but I left a TODO to address this more thoroughly in a follow up.
| previous_episode, previous_step = episode, step | ||
| native.save_episode(parallel_encoding=False) | ||
| native.finalize() | ||
| except BaseException: |
There was a problem hiding this comment.
Wouldn't this leave the root with a partial dataset that is marked as valid? can we use rmtree to delete the partial dataset and re-raise an error here in that scenario?
There was a problem hiding this comment.
Fixed, thanks for the callout
|
@bsprenger build-docs ci was triggered before i merged the fix in. |
github-actions
Bot
added
the
CI
2da2b1c to
027d7e4
Compare
Description
This PR adds a LeRobot replay dataset adapter, which can be used to export a replay buffer's contents to the LeRobot format.
Motivation and Context
LeRobot is a widely used robotics dataset format. Conversion to/from LeRobot datasets will allow sharing of trajectories for offline learning, and for pre-seeding replay buffers with LeRobot-format data from the Hugging Face hub.
An upcoming PR will add examples and benchmarks.
Types of changes
What types of changes does your code introduce? Remove all that do not apply:
Checklist
Go over all the following points, and put an
xin all the boxes that apply.If you are unsure about any of these, don't hesitate to ask. We are here to help!